"""Add name column to input artifact PK [728c6369cfaa].

Revision ID: 728c6369cfaa
Revises: 0.36.0
Create Date: 2023-03-20 13:37:51.215760

"""

from alembic import op
from sqlalchemy import text

# revision identifiers, used by Alembic.
revision = "728c6369cfaa"
down_revision = "0.36.0"
branch_labels = None
depends_on = None


def upgrade() -> None:
    """Upgrade database schema and/or data, creating a new revision.

    Raises:
        NotImplementedError: If the database engine is not SQLite or MySQL.
    """
    # ### commands auto generated by Alembic - please adjust! ###
    _disable_primary_key_requirement_if_necessary()

    engine_name = op.get_bind().engine.name
    if engine_name == "sqlite":
        constraint_name = "pk_step_run_input_artifact"
    elif engine_name == "mysql":
        constraint_name = "PRIMARY"
    else:
        raise NotImplementedError(f"Unsupported engine: {engine_name}")

    with op.batch_alter_table(
        "step_run_input_artifact",
        schema=None,
        naming_convention={"pk": "pk_%(table_name)s"},
    ) as batch_op:
        # Need to first drop the foreign keys as otherwise the table renaming
        # used during the PK constraint drop/create doesn't work
        batch_op.drop_constraint(
            constraint_name="fk_step_run_input_artifact_step_id_step_run",
            type_="foreignkey",
        )
        batch_op.drop_constraint(
            constraint_name="fk_step_run_input_artifact_artifact_id_artifact",
            type_="foreignkey",
        )

        # Update the PK
        batch_op.drop_constraint(
            constraint_name=constraint_name, type_="primary"
        )
        batch_op.create_primary_key(
            constraint_name="pk_step_run_input_artifact",
            columns=["step_id", "artifact_id", "name"],
        )

        # Re-add the foreign keys
        batch_op.create_foreign_key(
            constraint_name="fk_step_run_input_artifact_step_id_step_run",
            referent_table="step_run",
            local_cols=["step_id"],
            remote_cols=["id"],
            ondelete="CASCADE",
        )
        batch_op.create_foreign_key(
            constraint_name="fk_step_run_input_artifact_artifact_id_artifact",
            referent_table="artifact",
            local_cols=["artifact_id"],
            remote_cols=["id"],
            ondelete="CASCADE",
        )
    # ### end Alembic commands ###


def downgrade() -> None:
    """Downgrade database schema and/or data back to the previous revision."""
    # ### commands auto generated by Alembic - please adjust! ###

    _disable_primary_key_requirement_if_necessary()

    with op.batch_alter_table(
        "step_run_input_artifact",
        schema=None,
    ) as batch_op:
        # Need to first drop the foreign keys as otherwise the table renaming
        # used during the PK constraint drop/create doesn't work
        batch_op.drop_constraint(
            constraint_name="fk_step_run_input_artifact_step_id_step_run",
            type_="foreignkey",
        )
        batch_op.drop_constraint(
            constraint_name="fk_step_run_input_artifact_artifact_id_artifact",
            type_="foreignkey",
        )

        # Update the PK
        batch_op.drop_constraint(
            constraint_name="pk_step_run_input_artifact", type_="primary"
        )
        batch_op.create_primary_key(
            constraint_name="pk_step_run_input_artifact",
            columns=["step_id", "artifact_id"],
        )

        # Re-add the foreign keys
        batch_op.create_foreign_key(
            constraint_name="fk_step_run_input_artifact_step_id_step_run",
            referent_table="step_run",
            local_cols=["step_id"],
            remote_cols=["id"],
            ondelete="CASCADE",
        )
        batch_op.create_foreign_key(
            constraint_name="fk_step_run_input_artifact_artifact_id_artifact",
            referent_table="artifact",
            local_cols=["artifact_id"],
            remote_cols=["id"],
            ondelete="CASCADE",
        )

    # ### end Alembic commands ###


def _disable_primary_key_requirement_if_necessary() -> None:
    """Adjusts settings based on database engine requirements.

    Raises:
        NotImplementedError: If the database engine is not MySQL or mariadb.
    """
    conn = op.get_bind()
    engine = conn.engine
    engine_name = engine.name.lower()

    if engine_name == "mysql":
        server_version_info = engine.dialect.server_version_info
        if server_version_info is not None and server_version_info >= (
            8,
            0,
            13,
        ):
            potential_session_var = conn.execute(
                text('SHOW SESSION VARIABLES LIKE "sql_require_primary_key";')
            ).fetchone()
            if potential_session_var and potential_session_var[1] == "ON":
                # Temporarily disable this MySQL setting
                # for primary key modification
                op.execute("SET SESSION sql_require_primary_key = 0;")
    elif engine_name == "mariadb":
        # mariadb doesn't support session-based scoping for this setting
        op.execute("SET GLOBAL innodb_force_primary_key = 0;")
    else:
        raise NotImplementedError(f"Unsupported engine: {engine_name}")
